#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Elasticsearch 重索引工具

用于对指定的 ES 索引列表进行重索引操作，主要解决索引字段类型不符合要求的问题。
"""

import argparse
import json
import logging
import sys
import time
from typing import List, Dict, Any

from elasticsearch import Elasticsearch
from elasticsearch.exceptions import NotFoundError, ConnectionError


class ESReindexTool:
    """Elasticsearch 重索引工具类"""
    
    def __init__(self, es_host: str, es_port: int = 9200, 
                 username: str = None, password: str = None, 
                 use_ssl: bool = False, verify_certs: bool = True,
                 silent: bool = False):
        """
        初始化 ES 连接
        
        Args:
            es_host: ES 主机地址
            es_port: ES 端口，默认 9200
            username: 用户名（可选）
            password: 密码（可选）
            use_ssl: 是否使用 SSL，默认 False
            verify_certs: 是否验证证书，默认 True
            silent: 静默模式，不进行用户确认，默认 False
        """
        self.logger = self._setup_logger()
        self.silent = silent
        
        # 构建 ES 连接配置
        # 确定协议方案
        scheme = 'https' if use_ssl else 'http'
        
        # 构建主机 URL
        host_url = f"{scheme}://{es_host}:{es_port}"
        
        # 兼容新旧版本的连接配置
        try:
            # 尝试新版本的连接方式 (elasticsearch >= 8.0)
            es_config = {
                'hosts': [host_url],
                'timeout': 60,
                'max_retries': 3,
                'retry_on_timeout': True
            }
            
            if username and password:
                es_config['basic_auth'] = (username, password)
                
            if use_ssl and not verify_certs:
                es_config['verify_certs'] = False
                es_config['ssl_show_warn'] = False
                
            self.es = Elasticsearch(**es_config)
            
        except Exception as e:
            # 如果新版本方式失败，尝试旧版本方式 (elasticsearch < 8.0)
            self.logger.warning(f"尝试新版本连接方式失败，使用旧版本方式: {e}")
            
            es_config = {
                'hosts': [{'host': es_host, 'port': es_port, 'scheme': scheme}],
                'timeout': 60,
                'max_retries': 3,
                'retry_on_timeout': True
            }
            
            if username and password:
                es_config['http_auth'] = (username, password)
                
            if use_ssl:
                es_config['use_ssl'] = True
                es_config['verify_certs'] = verify_certs
                
            self.es = Elasticsearch(**es_config)
            
        try:
            # 测试连接
            info = self.es.info()
            self.logger.info(f"成功连接到 Elasticsearch {info['version']['number']}")
        except ConnectionError as e:
            self.logger.error(f"连接 Elasticsearch 失败: {e}")
            sys.exit(1)
    
    def _setup_logger(self) -> logging.Logger:
        """设置日志记录器"""
        logger = logging.getLogger('es_reindex')
        logger.setLevel(logging.INFO)
        
        if not logger.handlers:
            handler = logging.StreamHandler()
            formatter = logging.Formatter(
                '%(asctime)s - %(name)s - %(levelname)s - %(message)s'
            )
            handler.setFormatter(formatter)
            logger.addHandler(handler)
            
        return logger
    
    def check_index_exists(self, index_name: str) -> bool:
        """检查索引是否存在"""
        try:
            return self.es.indices.exists(index=index_name)
        except Exception as e:
            self.logger.error(f"检查索引 {index_name} 是否存在时出错: {e}")
            return False
    
    def get_index_mapping(self, index_name: str) -> Dict[str, Any]:
        """获取索引的映射配置"""
        try:
            return self.es.indices.get_mapping(index=index_name)
        except Exception as e:
            self.logger.error(f"获取索引 {index_name} 映射配置失败: {e}")
            return {}
    
    def get_index_settings(self, index_name: str) -> Dict[str, Any]:
        """获取索引的设置配置"""
        try:
            return self.es.indices.get_settings(index=index_name)
        except Exception as e:
            self.logger.error(f"获取索引 {index_name} 设置配置失败: {e}")
            return {}
    
    def create_temp_index(self, source_index: str) -> str:
        """创建临时索引"""
        temp_index = f"{source_index}_temp_{int(time.time())}"
        
        try:
            # 获取源索引的映射和设置
            mapping = self.get_index_mapping(source_index)
            settings = self.get_index_settings(source_index)
            
            # 提取源索引的映射和设置
            source_mapping = mapping.get(source_index, {}).get('mappings', {})
            source_settings = settings.get(source_index, {}).get('settings', {})
            
            # 清理设置中的系统字段
            if 'index' in source_settings:
                source_settings['index'].pop('creation_date', None)
                source_settings['index'].pop('uuid', None)
                source_settings['index'].pop('version', None)
                source_settings['index'].pop('provided_name', None)
            
            # 创建临时索引
            body = {}
            if source_mapping:
                body['mappings'] = source_mapping
            if source_settings:
                body['settings'] = source_settings
                
            self.es.indices.create(index=temp_index, body=body)
            self.logger.info(f"成功创建临时索引: {temp_index}")
            return temp_index
            
        except Exception as e:
            self.logger.error(f"创建临时索引失败: {e}")
            raise
    
    def create_new_index(self, index_name: str) -> bool:
        """
        重新创建索引（应用最新的索引模板）
        
        这个方法会创建一个空的索引，让 ES 自动应用匹配的索引模板
        """
        try:
            # 创建空索引，让 ES 应用索引模板
            self.es.indices.create(index=index_name)
            self.logger.info(f"成功重新创建索引: {index_name} (应用了最新的索引模板)")
            return True
        except Exception as e:
            self.logger.error(f"重新创建索引 {index_name} 失败: {e}")
            return False
    
    def reindex_data(self, source_index: str, dest_index: str) -> bool:
        """执行重索引操作"""
        try:
            reindex_body = {
                "source": {"index": source_index},
                "dest": {"index": dest_index}
            }
            
            self.logger.info(f"开始重索引: {source_index} -> {dest_index}")
            
            # 执行重索引
            response = self.es.reindex(body=reindex_body, wait_for_completion=True)
            
            if response.get('failures'):
                self.logger.error(f"重索引过程中出现失败: {response['failures']}")
                return False
            
            total = response.get('total', 0)
            created = response.get('created', 0)
            
            self.logger.info(f"重索引完成: 总计 {total} 条记录，创建 {created} 条记录")
            return True
            
        except Exception as e:
            self.logger.error(f"重索引操作失败: {e}")
            return False
    
    def delete_index(self, index_name: str, confirm_message: str = None) -> bool:
        """
        删除索引
        
        Args:
            index_name: 要删除的索引名称
            confirm_message: 确认消息（可选）
        """
        # 如果不是静默模式，则需要用户确认
        if not self.silent:
            if confirm_message:
                print(f"\n警告: {confirm_message}")
            
            confirm = input(f"确认删除索引 '{index_name}' 吗？此操作不可逆！(y/N): ")
            if confirm.lower() not in ('y', 'yes'):
                self.logger.info(f"用户取消删除索引: {index_name}")
                return False
        
        try:
            self.es.indices.delete(index=index_name)
            self.logger.info(f"成功删除索引: {index_name}")
            return True
        except NotFoundError:
            self.logger.warning(f"索引 {index_name} 不存在，跳过删除")
            return True
        except Exception as e:
            self.logger.error(f"删除索引 {index_name} 失败: {e}")
            return False
    
    def reindex_single_index(self, index_name: str) -> bool:
        """
        对单个索引执行重索引操作
        
        步骤：
        1. 创建临时索引（复制源索引的映射和设置）
        2. 将源索引数据重索引到临时索引
        3. 删除源索引
        4. 重新创建源索引（应用最新的索引模板）
        5. 将临时索引数据重索引回源索引
        6. 删除临时索引
        """
        self.logger.info(f"开始处理索引: {index_name}")
        
        # 检查源索引是否存在
        if not self.check_index_exists(index_name):
            self.logger.error(f"源索引 {index_name} 不存在")
            return False
        
        temp_index = None
        try:
            # 步骤1: 创建临时索引
            temp_index = self.create_temp_index(index_name)
            
            # 步骤2: 将源索引重索引到临时索引
            if not self.reindex_data(index_name, temp_index):
                self.logger.error(f"将 {index_name} 重索引到 {temp_index} 失败")
                return False
            
            # 步骤3: 删除源索引
            if not self.delete_index(index_name, f"即将删除源索引 '{index_name}' 以便应用最新的索引模板"):
                self.logger.error(f"删除源索引 {index_name} 失败")
                return False
            
            # 步骤4: 重新创建源索引（应用最新的索引模板）
            if not self.create_new_index(index_name):
                self.logger.error(f"重新创建源索引 {index_name} 失败")
                return False
            
            # 步骤5: 将临时索引数据重索引回源索引
            if not self.reindex_data(temp_index, index_name):
                self.logger.error(f"将 {temp_index} 重索引回 {index_name} 失败")
                return False
            
            # 步骤6: 删除临时索引
            if not self.delete_index(temp_index):
                self.logger.warning(f"删除临时索引 {temp_index} 失败，请手动清理")
            
            self.logger.info(f"索引 {index_name} 重索引完成")
            return True
            
        except Exception as e:
            self.logger.error(f"处理索引 {index_name} 时发生异常: {e}")
            
            # 清理临时索引
            if temp_index and self.check_index_exists(temp_index):
                self.delete_index(temp_index)
                
            return False
    
    def reindex_indices(self, index_list: List[str]) -> Dict[str, bool]:
        """
        批量重索引多个索引
        
        Args:
            index_list: 索引名称列表
            
        Returns:
            包含每个索引处理结果的字典
        """
        results = {}
        
        self.logger.info(f"开始批量重索引，共 {len(index_list)} 个索引")
        
        for i, index_name in enumerate(index_list, 1):
            self.logger.info(f"处理进度: {i}/{len(index_list)}")
            results[index_name] = self.reindex_single_index(index_name)
            
            # 在索引之间添加短暂延迟，避免对 ES 造成过大压力
            if i < len(index_list):
                time.sleep(1)
        
        # 输出总结
        success_count = sum(1 for success in results.values() if success)
        fail_count = len(index_list) - success_count
        
        self.logger.info(f"批量重索引完成: 成功 {success_count} 个，失败 {fail_count} 个")
        
        if fail_count > 0:
            failed_indices = [idx for idx, success in results.items() if not success]
            self.logger.error(f"失败的索引: {', '.join(failed_indices)}")
        
        return results


def parse_arguments():
    """解析命令行参数"""
    parser = argparse.ArgumentParser(description='Elasticsearch 重索引工具')
    
    # ES 连接参数
    parser.add_argument('--host', required=True, help='Elasticsearch 主机地址')
    parser.add_argument('--port', type=int, default=9200, help='Elasticsearch 端口 (默认: 9200)')
    parser.add_argument('--username', help='用户名')
    parser.add_argument('--password', help='密码')
    parser.add_argument('--ssl', action='store_true', help='使用 SSL 连接')
    parser.add_argument('--no-verify-certs', action='store_true', help='不验证 SSL 证书')
    
    # 索引参数
    parser.add_argument('--indices', nargs='+', help='要重索引的索引名称列表')
    parser.add_argument('--indices-file', help='包含索引名称列表的文件路径（每行一个索引名）')
    
    # 其他参数
    parser.add_argument('--dry-run', action='store_true', help='仅显示将要处理的索引，不执行实际操作')
    parser.add_argument('--verbose', '-v', action='store_true', help='显示详细日志')
    parser.add_argument('--silent', '--yes', '-y', action='store_true', help='静默模式，不进行确认提示，自动同意所有删除操作')
    
    return parser.parse_args()


def load_indices_from_file(file_path: str) -> List[str]:
    """从文件加载索引列表"""
    try:
        with open(file_path, 'r', encoding='utf-8') as f:
            indices = [line.strip() for line in f if line.strip() and not line.startswith('#')]
        return indices
    except Exception as e:
        print(f"读取索引文件失败: {e}")
        sys.exit(1)


def main():
    """主函数"""
    args = parse_arguments()
    
    # 设置日志级别
    if args.verbose:
        logging.getLogger('es_reindex').setLevel(logging.DEBUG)
    
    # 获取索引列表
    indices = []
    if args.indices:
        indices = args.indices
    elif args.indices_file:
        indices = load_indices_from_file(args.indices_file)
    else:
        print("错误: 请指定 --indices 或 --indices-file 参数")
        sys.exit(1)
    
    if not indices:
        print("错误: 没有找到要处理的索引")
        sys.exit(1)
    
    print(f"将要处理的索引: {', '.join(indices)}")
    
    if args.dry_run:
        print("--dry-run 模式，不执行实际操作")
        return
    
    # 确认操作（除非是静默模式）
    if not args.silent:
        confirm = input("确认要执行重索引操作吗？这将会删除并重新创建指定的索引 (y/N): ")
        if confirm.lower() not in ('y', 'yes'):
            print("操作已取消")
            return
    else:
        print("静默模式已启用，将自动执行重索引操作")
    
    # 创建重索引工具实例
    try:
        reindex_tool = ESReindexTool(
            es_host=args.host,
            es_port=args.port,
            username=args.username,
            password=args.password,
            use_ssl=args.ssl,
            verify_certs=not args.no_verify_certs,
            silent=args.silent
        )
    except Exception as e:
        print(f"初始化重索引工具失败: {e}")
        sys.exit(1)
    
    # 执行重索引
    results = reindex_tool.reindex_indices(indices)
    
    # 输出最终结果
    print("\n=== 重索引结果 ===")
    for index_name, success in results.items():
        status = "成功" if success else "失败"
        print(f"{index_name}: {status}")
    
    # 如果有失败的索引，以非零状态码退出
    if not all(results.values()):
        sys.exit(1)


if __name__ == '__main__':
    main()
